#include <pmm.h>
#include <list.h>
#include <string.h>
#include <buddy_pmm.h>

#define MAX_BUDDY_ORDER 20
extern size_t first_ppn;

typedef struct 
{
    unsigned int max_order;
    list_entry_t free_array[MAX_BUDDY_ORDER + 1];
    unsigned int nr_free;
} free_buddy_t;

free_buddy_t buddy_s;

#define buddy_array (buddy_s.free_array)
#define max_order (buddy_s.max_order)
#define nr_free (buddy_s.nr_free)

static bool IsPowerOf2(size_t n)
{
    return n&(n-1)==0;
}

static unsigned int GetOrder(size_t n)
{
    int order=0;
    while(n>>1)
    {
        n>>=1;
        order++;
    }
    return order;
}

static unsigned FixUp(unsigned size) {
    if(IsPowerOf2(size))
    {
        return size;
    }
    size |= size >> 1;
    size |= size >> 2;
    size |= size >> 4;
    size |= size >> 8;
    size |= size >> 16;
    return size+1;
}

static unsigned FixDown(unsigned size) 
{
    if(IsPowerOf2(size))
    {
        return size;
    }
    return (FixUp(size)>>1);
}

static void show_buddy_array(void) 
{
    cprintf("[*]Core: Print buddy array:\n");
    for (int i = 0;i < max_order + 1;i ++) 
    {
        cprintf("%d order: ", i);
        list_entry_t *le = &(buddy_array[i]);
        while ((le = list_next(le)) != &(buddy_array[i])) 
        {
            struct Page *p = le2page(le, page_link);
            cprintf("%d ", 1 << (p->property));
        }
        cprintf("\n");
    }
    cprintf("---------------------------\n");
    return;
}

static void buddy_init()
{
    for(int i=0;i<MAX_BUDDY_ORDER;i++)
    {
        list_init(&buddy_array[i]);
    }
    max_order=0;
    nr_free=0;
    return;
}

static void buddy_init_memmap(struct Page* base,size_t n)
{
    assert(n>0);
    size_t pcnt=FixDown(n);
    size_t order=GetOrder(pcnt);
    struct Page* p=base;
    cprintf("[*]Core: AVA Page num after rounding down to powers of 2: %d = 2^%d\n", pcnt, order);
    for(;p!=base+pcnt;p++)
    {
        assert(PageReserved(p));
        p->flags=0;
        p->property=-1;
        set_page_ref(p,0);
    }
    max_order=order;
    nr_free=pcnt;
    list_add(&(buddy_array[max_order]),&(base->page_link));
    base->property=max_order;
    return;
}

static struct Page* buddy_find(struct Page *page) 
{
    unsigned int order = page->property;
    unsigned int buddy_ppn = first_ppn + ((1 << order) ^ (page2ppn(page) - first_ppn));
    cprintf("[*]Core: Page NO.%d 's buddy page on order %d is: %d\n", page2ppn(page), order, buddy_ppn);
    if (buddy_ppn > page2ppn(page))
    {
        return page + (buddy_ppn - page2ppn(page));
    }
    else 
    {
        return page - (page2ppn(page) - buddy_ppn);
    }
}

static void buddy_split(size_t n) 
{
    assert(n > 0 && n <= max_order);
    assert(!list_empty(&(buddy_array[n])));
    cprintf("[*]Core: SPLITTING!\n");
    struct Page *page_a;
    struct Page *page_b;

    page_a = le2page(list_next(&(buddy_array[n])), page_link);
    page_b = page_a + (1 << (n - 1));
    page_a->property = n - 1;
    page_b->property = n - 1;

    list_del(list_next(&(buddy_array[n])));
    list_add(&(buddy_array[n-1]), &(page_a->page_link));
    list_add(&(page_a->page_link), &(page_b->page_link));

    return;
}

static struct Page *buddy_alloc_pages(size_t n) 
{
    assert(n > 0);
    if (n > nr_free) 
    {
        return NULL;
    }
    struct Page *page = NULL;
    size_t pnum = FixUp(n);
    size_t order = 0;
    order = GetOrder(pnum);
    cprintf("[*]Core: Allocating %d-->%d = 2^%d pages ...\n", n, pnum, order);
    cprintf("[*]COre: Buddy array before ALLOC:\n");
    //show_buddy_array();
    while(list_empty(&buddy_array[order]))
    {
        for(int i=order;i<max_order+1;i++)
        {
            if(!list_empty(&buddy_array[i]))
            {
                buddy_split(i);
                break;
            }
        }
    }
    page = le2page(list_next(&(buddy_array[order])), page_link);
    list_del(list_next(&(buddy_array[order])));
    SetPageProperty(page);
    cprintf("[*]Core: Buddy array after ALLOC NO.%d page:\n", page2ppn(page));
    show_buddy_array();
    nr_free -= pnum;
    cprintf("[*]Core: nr_free: %d\n", nr_free);
    return page;
}

static void buddy_free_pages(struct Page *base, size_t n) 
{
    assert(n > 0);
    unsigned int pnum = 1 << (base->property);
    assert(FixUp(n) == pnum);
    cprintf("[*]Core: Freeing NO.%d page leading %d pages block: \n", page2ppn(base), pnum);
    struct Page* left_block = base;
    struct Page* buddy = NULL;
    struct Page* tmp = NULL;

    buddy = buddy_find(left_block);
    list_add(&(buddy_array[left_block->property]), &(left_block->page_link));
    cprintf("[*]Core: add to list\n");
    //show_buddy_array();
    while (!PageProperty(buddy) && left_block->property < max_order) 
    {
        cprintf("[*]Core: Buddy free, MERGING!\n");
        if (left_block > buddy)//buddy is on the left
         {
            left_block->property = -1;
            ClearPageProperty(left_block);
            //redefine the "left"
            tmp = left_block;
            left_block = buddy;
            buddy = tmp;
        }
        list_del(&(left_block->page_link));    
        list_del(&(buddy->page_link));
        left_block->property += 1;
        list_add(&(buddy_array[left_block->property]), &(left_block->page_link));
        show_buddy_array();
        buddy = buddy_find(left_block);
    }
    cprintf("[*]Core: Buddy array after FREE:\n");
    ClearPageProperty(left_block);
    nr_free += pnum;
    show_buddy_array();
    
    cprintf("[*]Core: nr_free: %d\n", nr_free);
    return;
}

static size_t buddy_nr_free_pages(void) 
{
    return nr_free;
}

static void basic_check(void) 
{
    struct Page *p0, *p1, *p2;
    p0 = p1 = p2 = NULL;
    assert((p0 = alloc_page()) != NULL);
    assert((p1 = alloc_page()) != NULL);
    assert((p2 = alloc_page()) != NULL);
    free_page(p0);
    free_page(p1);
    free_page(p2);
    show_buddy_array();

    assert((p0 = alloc_pages(4)) != NULL);
    assert((p1 = alloc_pages(2)) != NULL);
    assert((p2 = alloc_pages(1)) != NULL);
    free_pages(p0, 4);
    free_pages(p1, 2);
    free_pages(p2, 1);
    show_buddy_array();

    assert((p0 = alloc_pages(3)) != NULL);
    assert((p1 = alloc_pages(3)) != NULL);
    free_pages(p0, 3);
    free_pages(p1, 3);

    show_buddy_array();
    cprintf("buddy_check success!");
}

static void buddy_check(void) 
{
    basic_check();
}

const struct pmm_manager buddy_pmm_manager = {
    .name = "buddy_pmm_manager",
    .init = buddy_init,
    .init_memmap = buddy_init_memmap,
    .alloc_pages = buddy_alloc_pages,
    .free_pages = buddy_free_pages,
    .nr_free_pages = buddy_nr_free_pages,
    .check = buddy_check,
};